import torch
import numpy as np

rand = torch.rand((5, 5))  # 元素在[0,1)之间

print("shape:", rand)
print("device:", rand.device)

tensor = torch.tensor([
    [1, 3],
    [1, 3],
    [1, 3]
], dtype=torch.float32).cuda()
print("shape:", tensor.mean(dim=1))

x = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(x[0, 1])  # 访问第一行第二个元素
print(x[:, 1])  # 访问第二列
print(x[1, :])  # 访问第二行
print(x[:, :2])  # 访问前两列
